from PIL import Image


# 图片格式统一化为256*256分辨率
def make_regular_image(img, size=(256, 256)):
    return img.resize(size).convert('RGB')


# 图像分为4*4块，每块分辨率为64*64
def split_image(img, part_size=(64, 64)):
    w, h = img.size
    pw, ph = part_size
    # 错误时触发异常
    assert w % pw == h % ph == 0

    return [img.crop((i, j, i+pw, j+ph)).copy() for i in range(0, w, pw) for j in range(0, h, ph)]


# 计算每个4*4小块的直方图相似度
def hist_similar(lh, rh):
    assert len(lh) == len(rh)
    return sum(1 - (0 if li == rr else float(abs(li - rr)) / max(li, rr)) for li, rr in zip(lh, rh)) / len(lh)


# 多次调用hist_similar计算整个图的相似性
def calc_similar(li, ri):
    # return hist_similar(li.histogram(), ri.histogram())
    return sum(
        hist_similar(l.histogram(), r.histogram())
        for l, r in zip(split_image(li), split_image(ri))) / 16.0


# 根据路径计算相似度
def calc_similar_by_path(lf, rf):
    li, ri = make_regular_image(Image.open(lf)), make_regular_image(Image.open(rf))
    return calc_similar(li, ri)


# 生成报告数据？？？（留）
def make_doc_data(lf, rf):
    li, ri = make_regular_image(Image.open(lf)), make_regular_image(
        Image.open(rf))
    li.save(lf + '_regular.png')
    ri.save(rf + '_regular.png')
    fd = open('stat.csv', 'w')
    fd.write('\n'.join(
        l + ',' + r
        for l, r in zip(map(str, li.histogram()), map(str, ri.histogram()))))
    # print >>fd, '\n'
    # fd.write(','.join(map(str, ri.histogram())))
    fd.close()
    from PIL import ImageDraw
    li = li.convert('RGB')
    draw = ImageDraw.Draw(li)
    for i in range(0, 256, 64):
        draw.line((0, i, 256, i), fill='#ff0000')
        draw.line((i, 0, i, 256), fill='#ff0000')
    li.save(lf + '_lines.png')


if __name__ == '__main__':
    # path = r'test/TEST%d/%d.JPG'
    for i in range(1, 7):
        print('test_case_%d: %.3f%%' % (i, calc_similar_by_path('C:/python_test/video_retrieval/test/TEST%d/%d.JPG' % (i, 1), 'C:/python_test/video_retrieval/test/TEST%d/%d.JPG' % (i, 2))*100))
    # make_doc_data('C:/python_test/video_retrieval/test/TEST4/1.JPG', 'C:/python_test/video_retrieval/test/TEST4/2.JPG')
